import os,cv2
import numpy as np
from math import ceil

from dataset.data_aug import *
from dataset.dataset import collate_fn

import torch
import torch.utils.data as torchdata
import torch.utils.data as data
from torch.autograd import Variable

from models.resnet import *

import matplotlib.pyplot as plt

class dataset_pred(data.Dataset):
    def __init__(self, path,label, transforms=None):
        self.paths = path
        self.labels = label
        self.transforms = transforms

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, item):
        img_path = self.paths[item]
        img =cv2.imread(img_path)
        img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
        if self.transforms is not None:
            img = self.transforms(img)
        label = self.labels[item]
        return torch.from_numpy(img).float(), label

#########   set the GPU   ###########
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

#########  the input path   ##########
image_path=['/media/hszc/model/zxy/table/data/test/239.jpg']

#########   the dataset    ###########
label_def={0:"该表有5个刻度关键点",
           1:"该表有6个刻度关键点",
           2:"该表有7个刻度关键点"}

test_transforms= Compose([
        ExpandBorder(size=(272,272),resize=True),
        RandomHflip(),
        Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

label=[1]
data_set = {}
data_set['test'] = dataset_pred(path=image_path,label=label,transforms=test_transforms)

data_loader = {}
data_loader['test'] = torchdata.DataLoader(data_set['test'], batch_size=1, num_workers=4,
                                           shuffle=False, pin_memory=True, collate_fn=collate_fn)

#########  choose the model ###########
model_name = 'resnet50-out'
model =resnet50(pretrained=True)
resume = '/media/hszc/model/zxy/table/predictor/model/resnet50/weights-10-210-[1.0000].pth'
# print('resuming finetune from %s'%resume)

model.avgpool = torch.nn.AdaptiveAvgPool2d(output_size=1)
a=model.fc.in_features
model.fc = torch.nn.Linear(model.fc.in_features,3)

model.load_state_dict(torch.load(resume),strict=False)
model = model.cuda()
model.eval()

##########   forward     ###########
test_size = ceil(len(data_set['test']) / data_loader['test'].batch_size)

for batch_cnt_test, data_test in enumerate(data_loader['test']):
    inputs, labels = data_test
    inputs = Variable(inputs.cuda())
    labels = Variable(torch.from_numpy(np.array(labels)).long().cuda())

    outputs,_ = model(inputs)
    if isinstance(outputs, list):
        outputs = (outputs[0]+outputs[1])/2
    props, preds = torch.max(outputs, 1)

    test_preds = preds.data.cpu().numpy()[0]
    test_props = props.data.cpu().numpy()[0]
    true_label = labels.data.cpu().numpy()[0]

    print(label_def[test_preds]+'!')

